import sqlglot
import json


class ConditionParser:
    def __init__(self):
        self.sql = None
        self.parsed_sql = None

    def _extract_conditions(self):
        supported_op = {
            'BETWEEN': sqlglot.expressions.Between,
            'EQ': sqlglot.expressions.EQ,
            'IN': sqlglot.expressions.In,
            'LIKE': sqlglot.expressions.Like,
            '<': sqlglot.expressions.LT,
            '>': sqlglot.expressions.GT,
            '<=': sqlglot.expressions.LTE,
            '>=': sqlglot.expressions.GTE,
            '!=': sqlglot.expressions.NEQ
        }

        parsed = self.parsed_sql[0]
        conditions = []

        # 首先找到Literal
        literals = list(parsed.find_all(sqlglot.expressions.Literal))
        for literal in literals:
            try:
                liter_parent = literal.parent

                for op_name, op_class in supported_op.items():
                    if isinstance(liter_parent, op_class):
                        # 获取该值对应的表名
                        table_name = liter_parent.this.table
                        # 获取该值对应的列名
                        column_name = liter_parent.this.this.this
                        # 获取该表名-列名下的所有值
                        values_in_liter = list(liter_parent.find_all(sqlglot.expressions.Literal))
                        values = []
                        for value in values_in_liter:
                            values.append(value.this)
                        # 获取操作类型
                        condition = {'table_name_alias': table_name,
                                     'table_name_origin': '',   # 暂时先留空
                                     'column_name': column_name,
                                     'operator': op_name,
                                     'value': values}
                        conditions.append(condition)
            except Exception as e:
                pass

        unique_dicts = []
        seen = set()
        for d in conditions:
            dict_str = json.dumps(d, sort_keys=True)
            if dict_str not in seen:
                unique_dicts.append(d)
                seen.add(dict_str)
        return unique_dicts

    def _get_alias_map(self):
        parsed = self.parsed_sql[0]

        tables_map = {}

        # 找到所有表
        tables = parsed.find_all(sqlglot.expressions.Table)
        for table in tables:
            tables_map[table.alias] = table.this.this
        return tables_map

    def get_conditions(self, sql) -> list:
        """
        主函数，将SQL拆解为条件并返回
        Args:
            sql: 原始 SQL 查询语句
        Returns:
            conditions: 条件字典列表，包含表别名和表名信息
        """
        self.sql = sql
        self.parsed_sql = sqlglot.parse(sql)

        # 获取表别名映射
        alias_map = self._get_alias_map()

        # 提取所有条件
        conditions = self._extract_conditions()

        # 遍历 conditions 并填充表名和别名
        for condition in conditions:
            # 尝试将别名映射到原始表名
            table_name_alias = condition.get('table_name_alias', '')
            table_name_origin = alias_map.get(table_name_alias, '')

            # 设置表的原始名称
            condition['table_name_origin'] = table_name_origin or table_name_alias

            # 如果原始表名为空，设置为别名
            if not condition['table_name_origin']:
                condition['table_name_origin'] = table_name_alias

            # 同样地，如果别名为空，则使用原始表名
            if not table_name_alias:
                condition['table_name_alias'] = condition['table_name_origin']

        return conditions


if __name__ == '__main__':
    sql = """
    SELECT u.user_id, u.name, o.order_id, o.order_date, oi.product_id, p.product_name, oi.quantity, oi.quantity * p.price AS total_price
FROM users u
JOIN orders o ON u.user_id = o.user_id
JOIN order_items oi ON o.order_id = oi.order_id
JOIN products p ON oi.product_id = p.product_id
JOIN cities c ON u.city_id = c.city_id
WHERE o.order_date BETWEEN '2024-03-01' AND '2024-06-01'
  AND o.total_amount > 500
  AND c.city_name IN ('New York', 'Los Angeles')
  AND u.user_id IN (
    SELECT DISTINCT u2.user_id
    FROM users u2
    JOIN orders o2 ON u2.user_id = o2.user_id
    JOIN order_items oi2 ON o2.order_id = oi2.order_id
    JOIN products p2 ON oi2.product_id = p2.product_id
    WHERE oi2.quantity > 3
      AND p2.price BETWEEN 20 AND 200
  )
ORDER BY o.order_date DESC, u.name;

    """

    sql2 = """
    SELECT A FROM B WHERE C='x'
    """

    sql3 = """SELECT COUNT(*) 
FROM 航空公司航班信息表 
WHERE CAST(乘客数量 AS INTEGER) > 200;
"""
    cp = ConditionParser()
    conditions = cp.get_conditions(sql=sql)
    print(conditions)

